# Chat 模型推理方法示例：
'''
https://zhuanlan.zhihu.com/p/632287465
FP4 量化有硬件要求吗？
请注意，此方法仅与 GPU 兼容，因此无法在 CPU 上以 4 位量化模型。在 GPU 中，这种方法应该没有任何硬件要求，因此只要安装了 CUDA>=11.2，任何 GPU 都可以用于运行 4bit 量化。还要记住，计算不是在 4 位中完成的，权重和激活被压缩为该格式，并且计算仍然保持在所需的或本机 dtype 中。
'''
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig

#model_path = "baichuan-inc/Baichuan2-7B-Chat-4bits"
model_path = "Baichuan2-7B-Chat-4bits"
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
#GPU
# model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
#CPU
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32, trust_remote_code=True)

model.generation_config = GenerationConfig.from_pretrained(model_path)
messages = []
messages.append({"role": "user", "content": "解释一下“温故而知新”"})
response = model.chat(tokenizer, messages)
print(response)